import dagre from 'dagre';
import { type Node, type Edge, Position } from '@xyflow/react';
import { type NodeData } from '../types/behaviorTree';
import { LAYOUT_CONSTANTS } from '../constants/layout';

/**
 * @constant NODE_WIDTH
 * @description 行为树节点的默认宽度，用于布局计算。
 */
const NODE_WIDTH = LAYOUT_CONSTANTS.NODE_WIDTH;

/**
 * @constant NODE_HEIGHT
 * @description 行为树节点的默认高度，用于布局计算。
 */
const NODE_HEIGHT = LAYOUT_CONSTANTS.NODE_HEIGHT;

/**
 * @constant DEFAULT_RANK_SEP
 * @description Dagre 布局中，层级之间的默认间距。
 */
const DEFAULT_RANK_SEP = LAYOUT_CONSTANTS.DEFAULT_RANK_SEP;

/**
 * @constant DEFAULT_NODE_SEP
 * @description Dagre 布局中，同一层级节点之间的默认间距。
 */
const DEFAULT_NODE_SEP = LAYOUT_CONSTANTS.DEFAULT_NODE_SEP;

/**
 * @function getLayoutedElements
 * @description 使用 Dagre 算法对 React Flow 的节点和边进行自动布局。
 * @param {Node<NodeData>[]} nodes - React Flow 的节点数组。
 * @param {Edge[]} edges - React Flow 的边数组。
 * @param {boolean} [direction='TB'] - 布局方向，'TB' (从上到下) 或 'LR' (从左到右)。
 * @returns {Node<NodeData>[]} 布局后的节点数组，包含更新后的位置信息。
 */
export const getLayoutedElements = (
    nodes: Node<NodeData>[],
    edges: Edge[],
    direction: 'TB' | 'LR' = 'TB'
): Node<NodeData>[] => {
    const dagreGraph = new dagre.graphlib.Graph();
    dagreGraph.setDefaultEdgeLabel(() => ({}));

    // 设置布局方向
    dagreGraph.setGraph({ rankdir: direction, ranksep: DEFAULT_RANK_SEP, nodesep: DEFAULT_NODE_SEP });

    // 为每个节点设置 Dagre 属性（宽度和高度）
    nodes.forEach((node) => {
        dagreGraph.setNode(node.id, { width: NODE_WIDTH, height: NODE_HEIGHT });
    });

    // 为每条边设置 Dagre 属性
    edges.forEach((edge) => {
        dagreGraph.setEdge(edge.source, edge.target);
    });

    // 执行 Dagre 布局计算
    dagre.layout(dagreGraph);

    // 更新节点的位置信息
    return nodes.map((node) => {
        const nodeWithPosition = dagreGraph.node(node.id);
        
        // 计算新的节点位置，Dagre 的位置是节点的中心点，React Flow 需要左上角位置
        node.position = {
            x: nodeWithPosition.x - NODE_WIDTH / 2,
            y: nodeWithPosition.y - NODE_HEIGHT / 2,
        };

        // 设置节点的 source 和 target handle 位置，以适应布局方向
        node.targetPosition = direction === 'LR' ? Position.Left : Position.Top;
        node.sourcePosition = direction === 'LR' ? Position.Right : Position.Bottom;

        return node;
    });
};